package com.note.es.utils;

import com.alibaba.fastjson.JSONObject;
import com.note.es.connection.ESConnection;
import com.note.es.vo.*;
import org.elasticsearch.action.search.*;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.aggregations.Aggregations;
import org.elasticsearch.search.aggregations.bucket.histogram.ParsedDateHistogram;
import org.elasticsearch.search.aggregations.bucket.terms.Terms;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.StringUtils;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;

/**
 * Created by sangpeng on 2018/11/29.
 */
public class TwoAggreationUtil {

    private static final Logger logger = LoggerFactory.getLogger(TwoAggreationUtil.class);

    /**
     * 普通检索
     *
     * @param searchRequestVO
     * @return
     * @throws IOException
     */
    public SearchResponseVO search(SearchRequestVO searchRequestVO) throws IOException {
        SearchRequest searchRequest = new SearchRequest();
        searchRequest.indices(searchRequestVO.getIndexName());
        searchRequest.types(searchRequestVO.getType());
        searchRequest.searchType(SearchType.QUERY_THEN_FETCH);
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        searchSourceBuilder.query(searchRequestVO.getQueryBuilder());
        searchSourceBuilder.from(searchRequestVO.getStartFrom());
        searchSourceBuilder.size(searchRequestVO.getPageSize());
        if (!StringUtils.isEmpty(searchRequestVO.getSortField())) {
            searchSourceBuilder.sort(searchRequestVO.getSortField(), searchRequestVO.getSortOrder());
        }
        if (searchRequestVO.getIncludeFields().length > 0 || searchRequestVO.getExcludeFields().length > 0) {
            searchSourceBuilder.fetchSource(searchRequestVO.getIncludeFields(), searchRequestVO.getExcludeFields());
        }
        searchRequest.source(searchSourceBuilder);
        logger.info("检索索引：{},\n检索语句：{}", Arrays.toString(searchRequest.indices()), searchRequest.source().toString());
        System.out.println("检索索引：" + Arrays.toString(searchRequest.indices()) + "\n检索语句：" + searchRequest.source().toString());
        SearchResponse searchResponse = ESConnection.getClient().search(searchRequest, RequestOptions.DEFAULT);
        SearchHits searchHits = searchResponse.getHits();
        SearchResponseVO<Map<String, Object>> searchResponseVO = new SearchResponseVO<>();
        searchResponseVO.setTotalCount(searchHits.getTotalHits());
        List<Map<String, Object>> list = new ArrayList<>();
        searchHits.iterator().forEachRemaining(hit -> {
            list.add(hit.getSourceAsMap());
        });
        searchResponseVO.setData(list);
        return searchResponseVO;
    }

    /**
     * 滚动查询
     *
     * @param scrollSearchRequestVO
     * @return
     * @throws IOException
     */
    public ScrollSearchResponseVO scrollSearch(ScrollSearchRequestVO scrollSearchRequestVO) throws IOException {

        if (StringUtils.isEmpty(scrollSearchRequestVO.getScrollId())) {
            SearchRequest searchRequest = new SearchRequest();
            searchRequest.indices(scrollSearchRequestVO.getIndexName());
            searchRequest.types(scrollSearchRequestVO.getType());
//            searchRequest.searchType(SearchType.QUERY_THEN_FETCH);
            SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
            searchSourceBuilder.query(scrollSearchRequestVO.getQueryBuilder());
            searchSourceBuilder.size(scrollSearchRequestVO.getLimit());
            if (scrollSearchRequestVO.getIncludeFields().length > 0 || scrollSearchRequestVO.getExcludeFields().length > 0) {
                searchSourceBuilder.fetchSource(scrollSearchRequestVO.getIncludeFields(), scrollSearchRequestVO.getExcludeFields());
            }
            searchRequest.source(searchSourceBuilder);
            searchRequest.scroll(TimeValue.timeValueSeconds(scrollSearchRequestVO.getKeepAlive()));
            logger.info("检索索引：{},\n检索语句：{}", Arrays.toString(searchRequest.indices()), searchRequest.source().toString());
            System.out.println("检索索引：" + Arrays.toString(searchRequest.indices()) + "\n检索语句：" + searchRequest.source().toString());
            SearchResponse searchResponse = ESConnection.getClient().search(searchRequest, RequestOptions.DEFAULT);
            SearchHits searchHits = searchResponse.getHits();
            ScrollSearchResponseVO<Map<String, Object>> scrollSearchResponseVO = new ScrollSearchResponseVO<>();
            scrollSearchResponseVO.setScrollId(searchResponse.getScrollId());
            scrollSearchResponseVO.setTotalCount(searchHits.getTotalHits());
            List<Map<String, Object>> list = new ArrayList<>();
            searchHits.iterator().forEachRemaining(hit -> {
                list.add(hit.getSourceAsMap());
            });
            scrollSearchResponseVO.setData(list);
            return scrollSearchResponseVO;
        } else {
            return notFirstScrollSearch(scrollSearchRequestVO);
        }

    }

    /**
     * 非首次滚动查询
     *
     * @param scrollSearchRequestVO
     * @return
     * @throws IOException
     */
    private ScrollSearchResponseVO notFirstScrollSearch(ScrollSearchRequestVO scrollSearchRequestVO) throws IOException {
        SearchScrollRequest searchScrollRequest = new SearchScrollRequest(scrollSearchRequestVO.getScrollId());
        searchScrollRequest.scroll(TimeValue.timeValueSeconds(scrollSearchRequestVO.getKeepAlive()));
        SearchResponse searchResponse = ESConnection.getClient().scroll(searchScrollRequest, RequestOptions.DEFAULT);
        SearchHits searchHits = searchResponse.getHits();
        ScrollSearchResponseVO<Map<String, Object>> scrollSearchResponseVO = new ScrollSearchResponseVO<>();
        scrollSearchResponseVO.setScrollId(searchResponse.getScrollId());
        scrollSearchResponseVO.setTotalCount(searchHits.getTotalHits());
        List<Map<String, Object>> list = new ArrayList<>();
        searchHits.iterator().forEachRemaining(hit -> {
            list.add(hit.getSourceAsMap());
        });
        scrollSearchResponseVO.setData(list);
        return scrollSearchResponseVO;
    }


    /**
     * 清除滚动查询游标
     *
     * @param scrollId
     * @return
     * @throws IOException
     */
    public boolean clearScroll(String scrollId) throws IOException {
        ClearScrollRequest clearScrollRequest = new ClearScrollRequest();
        clearScrollRequest.addScrollId(scrollId);
        ClearScrollResponse clearScrollResponse = ESConnection.getClient()
                .clearScroll(clearScrollRequest, RequestOptions.DEFAULT);
        return clearScrollResponse.isSucceeded();
    }

    /**
     * 聚合（对一/二级聚合友好）
     *
     * @param aggRequestVO
     */
    public List<JSONObject> aggs(AggRequestVO aggRequestVO) throws IOException {
        SearchRequest searchRequest = new SearchRequest();
        searchRequest.indices(aggRequestVO.getIndexName());
        searchRequest.types(aggRequestVO.getType());
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        searchSourceBuilder.query(aggRequestVO.getQueryBuilder());
        searchSourceBuilder.aggregation(aggRequestVO.getAggregationBuilder());
        searchSourceBuilder.size(0);
        searchRequest.source(searchSourceBuilder);
        logger.info("检索索引：{},\n检索语句：{}", Arrays.toString(searchRequest.indices()), searchRequest.source().toString());
        System.out.println("检索索引：" + Arrays.toString(searchRequest.indices()) + "\n检索语句：" + searchRequest.source().toString());
        Aggregations aggregations = ESConnection.getClient().search(searchRequest, RequestOptions.DEFAULT)
                .getAggregations();
        List<JSONObject> list = new ArrayList<>();
        aggResultParser(aggregations, list, null);
        return list;
    }

    /**
     * 聚合结果解析
     *
     * @param aggregations
     * @param list
     * @param node
     */
    private void aggResultParser(Aggregations aggregations, List<JSONObject> list, JSONObject node) {
        aggregations.iterator().forEachRemaining(aggregation -> {
            Object obj = aggregations.get(aggregation.getName());
            if (obj instanceof Terms) {
                Terms terms = (Terms) obj;
                terms.getBuckets().forEach(bucket -> {
                    bucketParser(list, node, bucket.getKeyAsString(), bucket.getDocCount(), bucket.getAggregations());
                });
            } else if (obj instanceof ParsedDateHistogram) {
                ParsedDateHistogram terms = (ParsedDateHistogram) obj;
                terms.getBuckets().forEach(bucket -> {
                    bucketParser(list, node, bucket.getKeyAsString(), bucket.getDocCount(), bucket.getAggregations());
                });
            }
        });
    }

    /**
     * bucket 解析
     *
     * @param list
     * @param node
     * @param keyAsString
     * @param docCount
     * @param aggregations
     */
    private void bucketParser(List<JSONObject> list,
                              JSONObject node,
                              String keyAsString,
                              long docCount,
                              Aggregations aggregations) {
        if (node == null) {
            node = new JSONObject();
            JSONObject parentNode = new JSONObject();
            parentNode.put(keyAsString, docCount);
            node.put("node", parentNode);
            list.add(node);
        } else {
            JSONObject childNode = node.getJSONObject("childNode");
            if (childNode == null) {
                childNode = new JSONObject();
                node.put("childNode", childNode);
            }
            childNode.put(keyAsString, docCount);
        }
        if (!aggregations.asList().isEmpty()) {
            aggResultParser(aggregations, list, node);
        }
    }
}
